import torch
import torch.nn as nn
import torch.nn.functional as F
import json
import numpy as np

Data=json.load(open('Attack_select/Configs/Config_net.json','r',encoding='UTF-8'))
N_S,N_A,GAMMA,MAX_EP,Batch_Size,_,_=Data.values()

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.Seq = nn.Sequential(
            nn.Linear(N_S, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Linear(1024, N_A)
        )

    def forward(self,input):
        input = torch.Tensor(input).cuda()
        out = self.Seq(input)
        return out

    def loss(self,input,target):
        target=torch.LongTensor(target).cuda()
        loss_fn=nn.CrossEntropyLoss()
        loss=loss_fn(input,target)
        return loss

    def set_init(self,layers):   # 初始化网络
        for layer in layers:
            nn.init.normal_(layer.weight, mean=0., std=0.1)
            nn.init.constant_(layer.bias, 0.1)

class Net2(nn.Module):
    def __init__(self):
        N_S=12
        N_A=66
        super(Net2,self).__init__()
        self.Seq=nn.Sequential(
            nn.Linear(N_S,512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512,512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512,N_A)
            )

    def forward(self,input):
        input=torch.Tensor(input).cuda()
        out=self.Seq(input)
        return out

    def loss(self,input,target):
        target=torch.LongTensor(target).cuda()
        loss_fn=nn.CrossEntropyLoss()
        loss=loss_fn(input,target)
        return loss

    def set_init(self,layers):   # 初始化网络
        for layer in layers:
            nn.init.normal_(layer.weight, mean=0., std=0.1)
            nn.init.constant_(layer.bias, 0.1)

